import sys
import cv2
import mediapipe as mp
import numpy as np
import argparse
import logging
from typing import Tuple, Optional

# 正确定义姿态关键点模块
mp_pose = mp.solutions.pose  # 关键修正1

# 配置日志格式
logging.basicConfig(
    level=logging.INFO,
    format='[%(asctime)s] [%(levelname)s] %(message)s',
    handlers=[logging.StreamHandler(sys.stderr)]
)

class PoseAnalyzer:
    def __init__(self, complexity: int = 1):
        self.pose = mp_pose.Pose(  # 关键修正2: 使用mp_pose
            static_image_mode=False,
            model_complexity=complexity,
            min_detection_confidence=0.7,
            min_tracking_confidence=0.7
        )
        self.draw_utils = mp.solutions.drawing_utils
        logging.info("姿态检测模型初始化完成 (复杂度 %d)", complexity)

    def process_image(self, image: np.ndarray) -> Tuple[Optional[np.ndarray], str]:
        """处理单帧图像并返回结果"""
        try:
            if image is None or image.size == 0:
                logging.error("输入图像  return None", "无效输入")

            # 转换为RGB格式
            rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            
            # 姿态检测
            results = self.pose.process(rgb_image)
            
            # 绘制关键点
            annotated_image = image.copy()
            action = "未检测到人体"
            if results.pose_landmarks:
                self._draw_landmarks(annotated_image, results.pose_landmarks)
                action = self._analyze_pose(results.pose_landmarks.landmark)
                return annotated_image, action
            return annotated_image, action    
        except Exception as e:
            logging.error("处理异常: %s", str(e), exc_info=True)
            return None, "处理错误"

    def _draw_landmarks(self, image: np.ndarray, landmarks) -> None:
        """绘制骨架连接线"""
        self.draw_utils.draw_landmarks(
            image,
            landmarks,
            mp_pose.POSE_CONNECTIONS,  # 关键修正3
            landmark_drawing_spec=mp.solutions.drawing_styles.get_default_pose_landmarks_style()
        )

    def _analyze_pose(self, landmarks):
        """核心动作分析逻辑（修复版）"""
        # 正确定义关键点常量
        LEFT_SHOULDER = mp_pose.PoseLandmark.LEFT_SHOULDER  # 关键修正4
        RIGHT_SHOULDER = mp_pose.PoseLandmark.RIGHT_SHOULDER
        LEFT_WRIST = mp_pose.PoseLandmark.LEFT_WRIST
        RIGHT_WRIST = mp_pose.PoseLandmark.RIGHT_WRIST

        try:
            # 安全获取关键点
            left_shoulder = landmarks[LEFT_SHOULDER]
            right_shoulder = landmarks[RIGHT_SHOULDER]
            left_wrist = landmarks[LEFT_WRIST] 
            right_wrist = landmarks[RIGHT_WRIST]

            actions = []
            # 手臂抬起检测
            if left_wrist.y < left_shoulder.y:
                actions.append("左臂抬起")
            if right_wrist.y < right_shoulder.y:
                actions.append("右臂抬起")

            return " | ".join(actions) if actions else "正常姿势"
        except (IndexError, KeyError) as e:
            logging.warning("关键点索引错误: %s", str(e))
            return "关键点数据不完整"
        except AttributeError as e:
            logging.warning("关键点属性错误: %s", str(e))
            return "数据解析失败"

def process_image_file(analyzer: PoseAnalyzer, path: str):
    """修复后的图片处理逻辑"""
    image = cv2.imread(path)
    if image is None:
        logging.error("无法读取图片文件: %s", path)
        return
        
    processed_image, action = analyzer.process_image(image)
    
    # 关键修正5: 严格检查输出有效性
    if processed_image is not None and processed_image.size > 0:
        try:
            success = cv2.imwrite('output.jpg', processed_image)
            if success:
                logging.info("结果已保存为 output.jpg")
                logging.info("检测结果: %s", action)
            else:
                logging.error("图片保存失败")
        except Exception as e:
            logging("保存失败: %s", str(e))
    else:
        logging.warning("未生成有效输出图像")

# ... main()及其他函数保持不变 ...

def main():
    # 参数解析
    parser = argparse.ArgumentParser(description='人体姿态检测程序')
    parser.add_argument('--input', type=str, default='webcam', 
                      help='输入源: webcam, image.jpg, video.mp4')
    parser.add_argument('--complexity', type=int, default=1,
                      help='模型复杂度 (0-2)')
    args = parser.parse_args()

    analyzer = PoseAnalyzer(complexity=args.complexity)
    
    # 处理不同输入源
    if args.input == 'webcam':
        process_webcam(analyzer)
    elif args.input.endswith(('.jpg', '.png')):
        process_image_file(analyzer, args.input)
    elif args.input.endswith(('.mp4', '.avi')):
        process_video_file(analyzer, args.input)
    else:
        logging.error("不支持的输入类型")

def process_webcam(analyzer: PoseAnalyzer):
    """处理摄像头输入"""
    cap = cv2.VideoCapture(0)
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
            
        processed_frame, action = analyzer.process_image(frame)
        
        if processed_frame is not None:
            cv2.imshow('姿态检测', processed_frame)
            logging.info("检测结果: %s", action)
            
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
            
    cap.release()
    cv2.destroyAllWindows()

def process_image_file(analyzer: PoseAnalyzer, path: str):
    """处理单张图片"""
    image = cv2.imread(path)
    if image is None:
        logging.error("无法读取图片文件: %s", path)
        return
        
    processed_image, action = analyzer.process_image(image)
    
    if processed_image is not None:
        cv2.imshow('处理结果', processed_image)
        cv2.waitKey(0)
        cv2.destroyAllWindows()
        logging.info("最终检测结果: %s", action)
        
    cv2.imwrite('output.jpg', processed_image)
    logging.info("结果已保存为 output.jpg")

def process_video_file(analyzer: PoseAnalyzer, path: str):
    """处理视频文件"""
    cap = cv2.VideoCapture(path)
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
            
        processed_frame, action = analyzer.process_image(frame)
        
        if processed_frame is not None:
            cv2.imshow('视频检测', processed_frame)
            logging.info("当前帧结果: %s", action)
            
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
            
    cap.release()
    cv2.destroyAllWindows()

if __name__ == '__main__':
    main()